import os
import shutil

if __name__ == '__main__':
    # input_dir = '/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Generation/data/input'
    # output_dir = '/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Generation/data/output'
    error_dir = '/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Generation/archive/error'
    success_dir = '/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Generation/archive/success'
    input_dir = '/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Generation/data/input'
    output_dir = '/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Generation/data/output'

    results = os.listdir(output_dir)
    error_results = []
    success_results = []

    for result in results:
        output = os.path.join(output_dir, result)
        result_files = os.listdir(output)
        if 'result.pdb' not in result_files:
            error_results.append(result)
        else:
            success_results.append(result)
    
    for error_result in error_results:
        error_pdb = error_result[5:] + '.pdb'
        shutil.move(os.path.join(input_dir, error_pdb), os.path.join(error_dir, error_pdb))
        shutil.rmtree(os.path.join(output_dir, error_result))
    
    for success_result in success_results:
        success_file = success_result[5:]
        success_pdb = success_file + '.pdb'
        os.makedirs(os.path.join(success_dir, success_file), exist_ok=True)

        shutil.move(os.path.join(input_dir, success_pdb), os.path.join(success_dir, success_file, 'structure.pdb'))
        shutil.move(os.path.join(output_dir, success_result, 'result.pdb'), os.path.join(success_dir, success_file, 'dynamic.pdb'))
        shutil.rmtree(os.path.join(output_dir, success_result))
